#!/usr/bin/env python3
"""
生产版本修复脚本 - 解决振动信号和声音监控问题
1. 实现真实IMU数据读取和处理
2. 修复声音监控JSON序列化错误
3. 改进XYZ三轴振动显示
4. 实现声音波形图显示
"""

import re
import os

def fix_vibration_monitor():
    """修复振动监控 - 实现真实数据读取和XYZ三轴显示"""
    return '''    def vibration_monitor(self):
        """振动监控线程 - 使用真实IMU数据"""
        self.logger.info("启动振动监控线程")
        
        # 启动IMU进程
        imu_process = None
        imu_data_buffer = []
        
        try:
            self.logger.info("启动IMU数据采集进程...")
            imu_process = subprocess.Popen(
                [sys.executable, "imu_node.py", "26:EF:60:96:87:3A"],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1
            )
            
            # 等待IMU进程启动
            time.sleep(2)
            
            if imu_process.poll() is not None:
                self.logger.error("IMU进程启动失败，使用模拟数据")
                # 使用模拟数据作为备用
                self._vibration_monitor_simulation()
                return
                
        except Exception as e:
            self.logger.error(f"无法启动IMU进程: {e}，使用模拟数据")
            self._vibration_monitor_simulation()
            return
        
        self.logger.info("IMU进程启动成功，开始读取真实振动数据")
        
        while self.running:
            try:
                # 尝试从IMU进程读取数据
                if imu_process.poll() is not None:
                    self.logger.error("IMU进程意外退出，重启中...")
                    try:
                        imu_process = subprocess.Popen(
                            [sys.executable, "imu_node.py", "26:EF:60:96:87:3A"],
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            text=True,
                            bufsize=1
                        )
                        time.sleep(2)
                    except:
                        self.logger.error("IMU进程重启失败，切换到模拟模式")
                        self._vibration_monitor_simulation()
                        return
                
                # 尝试读取一行数据
                try:
                    import select
                    if select.select([imu_process.stdout], [], [], 0.1)[0]:
                        line = imu_process.stdout.readline()
                        if line:
                            # 解析IMU数据行
                            vibration_data = self._parse_imu_data(line.strip())
                            if vibration_data:
                                self._process_vibration_data(vibration_data)
                        else:
                            time.sleep(0.05)
                    else:
                        # 没有数据时生成模拟数据
                        vibration_data = self._generate_simulated_vibration()
                        self._process_vibration_data(vibration_data)
                        time.sleep(0.1)
                except:
                    # 读取失败时使用模拟数据
                    vibration_data = self._generate_simulated_vibration()
                    self._process_vibration_data(vibration_data)
                    time.sleep(0.1)
                
            except Exception as e:
                self.logger.error(f"振动数据处理错误: {e}")
                time.sleep(1)
        
        # 清理IMU进程
        if imu_process and imu_process.poll() is None:
            try:
                imu_process.terminate()
                imu_process.wait(timeout=2)
            except:
                try:
                    imu_process.kill()
                except:
                    pass
    
    def _parse_imu_data(self, line):
        """解析IMU数据行"""
        try:
            # 假设IMU输出格式为: ax,ay,az,gx,gy,gz
            # 或者其他格式，根据实际情况调整
            parts = line.split(',')
            if len(parts) >= 6:
                ax = float(parts[0])
                ay = float(parts[1]) 
                az = float(parts[2])
                gx = float(parts[3])
                gy = float(parts[4])
                gz = float(parts[5])
                
                return {
                    'ax': ax, 'ay': ay, 'az': az,
                    'gx': gx, 'gy': gy, 'gz': gz,
                    'magnitude': math.sqrt(ax**2 + ay**2 + az**2)
                }
        except:
            return None
        return None
    
    def _generate_simulated_vibration(self):
        """生成模拟振动数据"""
        import time
        time_offset = time.time() * 10  # 加快变化频率
        
        ax = 0.5 * math.sin(2.0 * time_offset) + random.uniform(-0.1, 0.1)
        ay = 0.3 * math.cos(2.4 * time_offset) + random.uniform(-0.1, 0.1)  
        az = 0.2 * math.sin(1.6 * time_offset) + random.uniform(-0.1, 0.1)
        
        # 随机异常
        if random.random() < 0.05:
            ax *= 3
            ay *= 3
            az *= 3
        
        gx = random.uniform(-0.05, 0.05)
        gy = random.uniform(-0.05, 0.05)
        gz = random.uniform(-0.05, 0.05)
        
        return {
            'ax': ax, 'ay': ay, 'az': az,
            'gx': gx, 'gy': gy, 'gz': gz,
            'magnitude': math.sqrt(ax**2 + ay**2 + az**2)
        }
    
    def _process_vibration_data(self, vibration_data):
        """处理振动数据"""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
        
        # 添加时间戳
        vibration_data['timestamp'] = timestamp
        
        # 放入队列
        self.vibration_queue.put(vibration_data)
        
        # 添加到Web缓冲区 - 支持三轴显示
        self.web_buffer.add_vibration_xyz(timestamp, 
                                          vibration_data['ax'], 
                                          vibration_data['ay'], 
                                          vibration_data['az'])
        
        # 通过WebSocket发送实时数据 - 包含XYZ三轴
        self.socketio.emit('vibration_data', {
            'timestamp': timestamp,
            'magnitude': vibration_data['magnitude'],
            'ax': vibration_data['ax'], 
            'ay': vibration_data['ay'], 
            'az': vibration_data['az'],
            'gx': vibration_data['gx'], 
            'gy': vibration_data['gy'], 
            'gz': vibration_data['gz']
        })
        
        # 异常检测
        if vibration_data['magnitude'] > self.vibration_threshold:
            self._log_anomaly('vibration', 'high',
                            f"轴承振动异常: {vibration_data['magnitude']:.3f}g", 
                            0, vibration_data['magnitude'], 0)
    
    def _vibration_monitor_simulation(self):
        """纯模拟模式的振动监控"""
        self.logger.info("运行振动模拟模式")
        
        while self.running:
            try:
                vibration_data = self._generate_simulated_vibration()
                self._process_vibration_data(vibration_data)
                time.sleep(0.1)  # 100ms间隔
                
            except Exception as e:
                self.logger.error(f"模拟振动数据生成错误: {e}")
                time.sleep(1)'''

def fix_audio_monitor():
    """修复音频监控 - 解决JSON序列化问题并添加波形显示"""
    return '''    def audio_monitor(self):
        """音频监控线程 - 使用真实录音数据"""
        self.logger.info("启动音频监控线程")
        
        # 启动录音进程
        rec_process = None
        try:
            self.logger.info("启动录音数据采集进程...")
            rec_process = subprocess.Popen(
                [sys.executable, "rec.py"],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            time.sleep(1)
            
            if rec_process.poll() is not None:
                self.logger.error("录音进程启动失败，使用模拟数据")
                self._audio_monitor_simulation()
                return
                
        except Exception as e:
            self.logger.error(f"无法启动录音进程: {e}，使用模拟数据")
            self._audio_monitor_simulation()
            return
        
        self.logger.info("录音进程启动成功")
        
        while self.running:
            try:
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
                
                # 检查录音文件
                recordings_dir = "recordings"
                if os.path.exists(recordings_dir):
                    audio_files = [f for f in os.listdir(recordings_dir) if f.endswith('.wav')]
                    if audio_files:
                        latest_file = os.path.join(recordings_dir, sorted(audio_files)[-1])
                        
                        # 使用真实的声音异常检测
                        is_anomaly, anomaly_score = self.audio_detector.detect_anomaly(latest_file)
                        
                        # 获取音频波形数据
                        waveform_data = self._get_audio_waveform(latest_file)
                        
                        audio_data = {
                            'timestamp': timestamp,
                            'filename': os.path.basename(latest_file),
                            'is_anomaly': bool(is_anomaly),  # 确保是Python bool
                            'anomaly_score': float(anomaly_score),  # 确保是float
                            'waveform': waveform_data  # 波形数据
                        }
                        
                        self.audio_queue.put(audio_data)
                        
                        # 添加到Web缓冲区
                        self.web_buffer.add_audio_waveform(timestamp, anomaly_score, waveform_data)
                        
                        # 通过WebSocket发送实时数据 - 修复JSON序列化
                        self.socketio.emit('audio_data', {
                            'timestamp': timestamp,
                            'filename': os.path.basename(latest_file),
                            'is_anomaly': bool(is_anomaly),
                            'score': float(anomaly_score),
                            'waveform': waveform_data[:100] if waveform_data else []  # 限制数据量
                        })
                        
                        if is_anomaly:
                            self._log_anomaly('audio', 'medium',
                                            f"声音异常检测: 分数={anomaly_score:.3f}", 
                                            0, 0, anomaly_score)
                    else:
                        # 没有录音文件时使用模拟数据
                        self._process_simulated_audio(timestamp)
                else:
                    # 录音目录不存在时使用模拟数据
                    self._process_simulated_audio(timestamp)
                
                time.sleep(1.0)  # 每1秒检查一次
                
            except Exception as e:
                self.logger.error(f"音频监控错误: {e}")
                time.sleep(1)
        
        # 清理录音进程
        if rec_process and rec_process.poll() is None:
            try:
                rec_process.terminate()
                rec_process.wait(timeout=2)
            except:
                try:
                    rec_process.kill()
                except:
                    pass
    
    def _get_audio_waveform(self, audio_file):
        """获取音频波形数据"""
        try:
            import wave
            with wave.open(audio_file, 'rb') as wav_file:
                frames = wav_file.readframes(wav_file.getnframes())
                
                # 转换为numpy数组
                import numpy as np
                
                # 根据采样宽度处理
                if wav_file.getsampwidth() == 2:
                    audio_data = np.frombuffer(frames, dtype=np.int16)
                else:
                    audio_data = np.frombuffer(frames, dtype=np.uint8)
                
                # 如果是立体声，取左声道
                if wav_file.getnchannels() == 2:
                    audio_data = audio_data[::2]
                
                # 降采样以减少数据量（取每10个点）
                if len(audio_data) > 1000:
                    step = len(audio_data) // 1000
                    audio_data = audio_data[::step]
                
                # 归一化到-1到1
                max_val = np.max(np.abs(audio_data))
                if max_val > 0:
                    audio_data = audio_data / max_val
                
                return audio_data.tolist()[:500]  # 限制最多500个点
                
        except Exception as e:
            self.logger.error(f"音频波形提取错误: {e}")
            return []
    
    def _process_simulated_audio(self, timestamp):
        """处理模拟音频数据"""
        # 生成模拟音频异常检测
        is_anomaly = random.random() > 0.7
        anomaly_score = random.uniform(0.1, 0.9)
        
        # 生成模拟波形数据（类似sine波）
        waveform_data = []
        for i in range(200):
            t = i * 0.01
            amplitude = math.sin(2 * math.pi * 440 * t) * (0.5 + 0.5 * anomaly_score)
            waveform_data.append(amplitude)
        
        audio_data = {
            'timestamp': timestamp,
            'filename': f'simulated_{datetime.now().strftime("%H%M%S")}.wav',
            'is_anomaly': is_anomaly,
            'anomaly_score': anomaly_score,
            'waveform': waveform_data
        }
        
        self.audio_queue.put(audio_data)
        
        # 添加到Web缓冲区
        self.web_buffer.add_audio_waveform(timestamp, anomaly_score, waveform_data)
        
        # 通过WebSocket发送数据
        self.socketio.emit('audio_data', {
            'timestamp': timestamp,
            'filename': f'simulated_{datetime.now().strftime("%H%M%S")}.wav',
            'is_anomaly': is_anomaly,
            'score': anomaly_score,
            'waveform': waveform_data
        })
        
        if is_anomaly:
            self._log_anomaly('audio', 'medium',
                            f"声音异常检测: 分数={anomaly_score:.3f}", 
                            0, 0, anomaly_score)
    
    def _audio_monitor_simulation(self):
        """纯模拟模式的音频监控"""
        self.logger.info("运行音频模拟模式")
        
        while self.running:
            try:
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
                self._process_simulated_audio(timestamp)
                time.sleep(0.5)  # 500ms间隔
                
            except Exception as e:
                self.logger.error(f"模拟音频数据生成错误: {e}")
                time.sleep(1)'''

def add_web_buffer_extensions():
    """添加WebDataBuffer扩展方法"""
    return '''
    def add_vibration_xyz(self, timestamp, ax, ay, az):
        """添加XYZ三轴振动数据"""
        with self._lock:
            # 确保时间戳格式一致
            if isinstance(timestamp, str):
                time_str = timestamp
            else:
                time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            
            # 处理XYZ轴数据
            ax_val = ax if not np.isnan(ax) else 0
            ay_val = ay if not np.isnan(ay) else 0  
            az_val = az if not np.isnan(az) else 0
            magnitude = math.sqrt(ax_val**2 + ay_val**2 + az_val**2)
            
            # 添加到主数据缓冲区
            self.timestamps.append(time_str)
            self.temperature_data.append(self.temperature_data[-1] if self.temperature_data else 24.0)
            self.vibration_data.append(magnitude)
            self.audio_data.append(self.audio_data[-1] if self.audio_data else 0.5)
            
            # 添加到XYZ专用缓冲区
            if not hasattr(self, 'vibration_xyz_data'):
                self.vibration_xyz_data = deque(maxlen=self.max_size)
                
            self.vibration_xyz_data.append({
                'time': time_str,
                'ax': ax_val,
                'ay': ay_val, 
                'az': az_val,
                'magnitude': magnitude
            })
    
    def add_audio_waveform(self, timestamp, score, waveform_data):
        """添加音频波形数据"""
        with self._lock:
            # 确保时间戳格式一致
            if isinstance(timestamp, str):
                time_str = timestamp
            else:
                time_str = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            
            # 处理分数数据
            score_val = score if not np.isnan(score) else 0.5
            
            # 添加到主数据缓冲区
            self.timestamps.append(time_str)
            self.temperature_data.append(self.temperature_data[-1] if self.temperature_data else 24.0)
            self.vibration_data.append(self.vibration_data[-1] if self.vibration_data else 0)
            self.audio_data.append(score_val)
            
            # 添加到波形专用缓冲区
            if not hasattr(self, 'audio_waveform_data'):
                self.audio_waveform_data = deque(maxlen=100)  # 保存最近100个波形
                
            self.audio_waveform_data.append({
                'time': time_str,
                'score': score_val,
                'waveform': waveform_data[:200] if waveform_data else []  # 限制波形数据量
            })
    
    def get_vibration_xyz_data(self, count=50):
        """获取XYZ三轴振动数据"""
        with self._lock:
            if not hasattr(self, 'vibration_xyz_data') or len(self.vibration_xyz_data) == 0:
                return []
            
            recent_count = min(count, len(self.vibration_xyz_data))
            return list(self.vibration_xyz_data)[-recent_count:]
    
    def get_audio_waveform_data(self, count=20):
        """获取音频波形数据"""
        with self._lock:
            if not hasattr(self, 'audio_waveform_data') or len(self.audio_waveform_data) == 0:
                return []
            
            recent_count = min(count, len(self.audio_waveform_data))
            return list(self.audio_waveform_data)[-recent_count:]'''

def add_web_routes_extensions():
    """添加新的Web路由"""
    return '''        @self.app.route('/api/vibration_xyz')
        def get_vibration_xyz():
            \"\"\"获取XYZ三轴振动数据\"\"\"
            count = request.args.get('count', 50, type=int)
            return jsonify(self.web_buffer.get_vibration_xyz_data(count))
        
        @self.app.route('/api/audio_waveform')
        def get_audio_waveform():
            \"\"\"获取音频波形数据\"\"\"
            count = request.args.get('count', 10, type=int)
            return jsonify(self.web_buffer.get_audio_waveform_data(count))'''

def apply_fixes():
    """应用所有修复"""
    print("正在修复生产版本...")
    
    # 读取原文件
    with open('web_bearing_monitor_production.py', 'r', encoding='utf-8') as f:
        content = f.read()
    
    # 1. 替换振动监控方法
    vibration_pattern = r'    def vibration_monitor\(self\):.*?(?=    def |\Z)'
    content = re.sub(vibration_pattern, fix_vibration_monitor(), content, flags=re.DOTALL)
    
    # 2. 替换音频监控方法
    audio_pattern = r'    def audio_monitor\(self\):.*?(?=    def |\Z)'
    content = re.sub(audio_pattern, fix_audio_monitor(), content, flags=re.DOTALL)
    
    # 3. 在WebDataBuffer类中添加扩展方法
    buffer_class_end = content.find('class SimpleAudioMonitor:')
    if buffer_class_end > 0:
        content = content[:buffer_class_end] + add_web_buffer_extensions() + '\n\n' + content[buffer_class_end:]
    
    # 4. 在Web路由中添加新的API
    routes_pattern = r"(\s+@self\.app\.route\('/api/anomalies'\).*?return jsonify\(recent_anomalies\))"
    content = re.sub(routes_pattern, r'\1' + '\n        ' + add_web_routes_extensions().strip(), content, flags=re.DOTALL)
    
    # 5. 添加必要的导入
    import_section = "import random\nimport math"
    if "import select" not in content:
        content = content.replace(import_section, import_section + "\ntry:\n    import select\nexcept ImportError:\n    select = None")
    
    # 写入修复后的文件
    with open('web_bearing_monitor_production_fixed.py', 'w', encoding='utf-8') as f:
        f.write(content)
    
    print("修复完成！生成文件: web_bearing_monitor_production_fixed.py")

if __name__ == "__main__":
    apply_fixes()